import sys
import requests
from PyQt5.QtWidgets import QApplication, QMainWindow, QSizePolicy, QVBoxLayout
from ui_main_window import Ui_MainWindow
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.figure import Figure
import pandas as pd
import os

# ... (其他 import 语句)
import matplotlib.pyplot as plt

# --- 解决中文乱码的关键代码 ---
plt.rcParams['font.sans-serif'] = ['SimHei']  # 指定默认字体为黑体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号'-'显示为方块的问题
# --- 代码结束 ---

# 服务器URL，与您第二周的server.py保持一致
SERVER_URL = "http://localhost:5000/get_consumption"

class MplCanvas(FigureCanvas):
    """一个通用的Matplotlib画布Widget"""
    def __init__(self, parent=None, width=5, height=4, dpi=100):
        fig = Figure(figsize=(width, height), dpi=dpi)
        self.axes = fig.add_subplot(111)
        super(MplCanvas, self).__init__(fig)
        self.setParent(parent)
        FigureCanvas.setSizePolicy(self, QSizePolicy.Expanding, QSizePolicy.Expanding)
        FigureCanvas.updateGeometry(self)

class EnergyApp(QMainWindow):
    # __init__ 是类的构造方法，必须在类的内部
    def __init__(self):
        super().__init__()
        self.ui = Ui_MainWindow()
        self.ui.setupUi(self)
        self.init_ui_controls()
        self.ui.query_button.clicked.connect(self.perform_query)
        self.setup_plots()

    # 所有下面的方法都应该在这个类里面，注意缩进
    def init_ui_controls(self):
        """初始化界面控件的内容"""
        for i in range(1, 13):
            self.ui.month_combo.addItem(f"{i:02d}")
        self.setWindowTitle("能耗查询客户端")

    def get_consumption_from_server(self, month, time):
        """从服务器获取能耗数据的函数"""
        params = {'month': str(month), 'time': str(time)}
        try:
            response = requests.get(SERVER_URL, params=params)
            if response.status_code == 200:
                return response.json()
            else:
                error_data = response.json()
                return {"error": f"服务器错误, 状态码 {response.status_code}: {error_data.get('error', '未知错误')}"}
        except requests.exceptions.ConnectionError:
            return {"error": "连接失败。请确保WSL中的服务器正在运行。"}
        except Exception as e:
            return {"error": f"发生未知错误: {str(e)}"}

    def perform_query(self):
        """执行查询并更新UI"""
        month_str = self.ui.month_combo.currentText()
        time_str = self.ui.time_input.text()
        if not time_str:
            self.ui.result_display.setText("错误：请输入时间。")
            return
        try:
            time_val = float(time_str)
            if not (0.0 <= time_val <= 23.5):
                self.ui.result_display.setText("错误：时间必须在 0.0 到 23.5 之间。")
                return
        except ValueError:
            self.ui.result_display.setText("错误：时间格式不正确，请输入数字。")
            return
        
        result = self.get_consumption_from_server(month_str, time_str)
        if "error" in result:
            display_text = f"查询失败: {result['error']}"
        else:
            display_text = (
                f"--- 服务器响应 ---\n"
                f"查询月份: {result['month']}\n"
                f"请求时间: {result['requested_time']}\n"
                f"对应小时: {result['hour_of_day']}\n"
                f"能耗: {result['energy_consumption']:.2f}"
            )
        self.ui.result_display.setText(display_text)

    def load_monthly_data(self, month_str):
        """加载并处理单月能耗数据，返回每小时平均值"""
        file_name = f'GenerationbyFuelType_2018{month_str}.csv'
        file_path = os.path.join('data', file_name)
        if not os.path.exists(file_path):
            print(f"警告: 文件 {file_path} 未找到")
            return None
        df = pd.read_csv(file_path, header=None)
        numeric_cols = df.columns[2:]
        df_numeric = df[numeric_cols].apply(pd.to_numeric, errors='coerce')
        half_hourly_sum = df_numeric.sum(axis=1)
        hourly_index = [i // 2 for i in range(len(half_hourly_sum))]
        hourly_sum = half_hourly_sum.groupby(hourly_index).sum()
        return hourly_sum

    def setup_plots(self):
        """设置绘图画布"""
        self.canvas1 = MplCanvas(self)
        layout1 = QVBoxLayout()
        layout1.addWidget(self.canvas1)
        self.ui.plot_widget_1.setLayout(layout1)

        self.canvas2 = MplCanvas(self)
        layout2 = QVBoxLayout()
        layout2.addWidget(self.canvas2)
        self.ui.plot_widget_2.setLayout(layout2)
        
        self.ui.month_combo.currentTextChanged.connect(self.update_plots)
        self.update_plots() # 初始加载一次图表

    def update_plots(self):
        """根据当前选择的月份更新两个图表"""
        selected_month_str = self.ui.month_combo.currentText()
        
        # 防止在UI初始化时因月份为空而报错
        if not selected_month_str:
            return
            
        selected_month_int = int(selected_month_str)
        current_month_data = self.load_monthly_data(selected_month_str)
        
        # --- 情况一：当月数据加载失败 ---
        if current_month_data is None:
            # 清空并提示折线图
            self.canvas1.axes.clear()
            self.canvas1.axes.set_title(f"月份 {selected_month_str} 数据缺失或路径错误")
            self.canvas1.draw()
            # 清空并提示饼图
            self.canvas2.axes.clear()
            self.canvas2.axes.set_title(f"月份 {selected_month_str} 数据缺失或路径错误")
            self.canvas2.draw()
            return

        # --- 情况二：数据加载成功，开始绘图 ---
        
        # 1. 准备前两个月的数据用于对比
        prev_months_data = []
        for i in range(1, 3):
            prev_month = selected_month_int - i
            if prev_month > 0:
                prev_month_str = f"{prev_month:02d}"
                data = self.load_monthly_data(prev_month_str)
                if data is not None:
                    prev_months_data.append(data)
        
        prev_two_month_avg = None
        if len(prev_months_data) > 0:
            prev_two_month_avg = pd.concat(prev_months_data, axis=1).mean(axis=1)

        # 2. 绘制折线图 (Canvas 1)
        self.canvas1.axes.clear()
        self.canvas1.axes.plot(current_month_data.index, current_month_data.values, label=f'{selected_month_str}月日均用电量')
        if prev_two_month_avg is not None:
            self.canvas1.axes.plot(prev_two_month_avg.index, prev_two_month_avg.values, label='前两月日均用电量', linestyle='--')
        
        self.canvas1.axes.set_title(f'{selected_month_str}月用电量分时折线图')
        self.canvas1.axes.set_xlabel('小时 (Hour)')
        self.canvas1.axes.set_ylabel('总能耗')
        
        # --- 关键修改：精确控制X轴 ---
        self.canvas1.axes.set_xticks(range(24))      # 设置刻度为 0, 1, ..., 23
        self.canvas1.axes.set_xlim(-0.5, 23.5)      # 设置X轴的显示范围
        # self.canvas1.axes.tick_params(axis='x', labelrotation=45) # (可选)如果觉得标签太挤，可以取消此行注释
        
        self.canvas1.axes.legend()
        self.canvas1.axes.grid(True)
        

        # 3. 准备饼图数据
        top_3 = current_month_data.nlargest(3)
        other_sum = current_month_data.sum() - top_3.sum()
        pie_data = top_3.copy()
        
        # 只有当除了top3外还有剩余时才添加“其他”项
        if other_sum > 0:
            pie_data['其他'] = other_sum
        
        # 4. 绘制饼图 (Canvas 2)
        self.canvas2.axes.clear()
        # 使用 f'{i}点' 作为标签，除了“其他”
        pie_labels = [f'{i}点' if isinstance(i, int) else i for i in pie_data.index]
        self.canvas2.axes.pie(pie_data, labels=pie_labels, autopct='%1.1f%%', startangle=90)
        
        self.canvas2.axes.set_title(f'{selected_month_str}月用电量分时饼图 (Top 3)')
        self.canvas2.axes.axis('equal')  # 保证饼图是正圆形

        # 5. 统一刷新两个画布
        self.canvas1.draw()
        self.canvas2.draw()

# 程序入口
if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = EnergyApp()
    window.show()
    sys.exit(app.exec_())